package lzc.encrypt.encryption;

import Jama.Matrix;
import it.unisa.dia.gas.jpbc.Element;
import it.unisa.dia.gas.jpbc.ElementPowPreProcessing;

import java.math.BigInteger;
import java.util.concurrent.ThreadLocalRandom;

/**
 * 密钥类，IPPE加密向量使用的密钥
 *
 * @author 蓝色的三明治
 * @date 2024/09/20
 */
public class EncryptionKey {
    //    private Matrix keyMatrix;// 基于 Z_p 的随机可逆矩阵，维度是 向量维度 + 2
    private BigIntMatrix keyMatrix;// 基于 Z_p 的随机可逆矩阵，维度是 向量维度 + 2
    private ElementPowPreProcessing number;// 用于加密的大整数，一般是 A 或 B，A = g^a，B = g^b，其中 a 和 b 是 Z^*_p 中的随机值
    private Element n;// 用于生成大整数的操作次数，即 a 或 b
    private boolean ifData;// true表示用于加密数据向量，false表示用于加密搜索向量

    /**
     * 初始化加密密钥组（生成随机可逆矩阵）
     *
     * @param g         循环群生成元
     * @param n         取自有效域的随机值
     * @param order     循环群的阶
     * @param dimension 向量维度
     */
    public EncryptionKey(Element g, Element n, BigInteger order, int dimension) {
        // 根据所给参数是向量的维度，确定该组密钥需要生成一个随机可逆矩阵，也就是说，这组密钥用于加密搜索向量
        ifData = false;

        // 记录属于 Z^*_p 的随机值 n，对生成元 g 进行 n 次运算
        this.n = n;
        number = g.powZn(n).getImmutable().getElementPowPreProcessing();

        // 生成一个随机可逆矩阵
//        keyMatrix = generateInvertibleMatrix(dimension + 2);
        keyMatrix = new BigIntMatrix(dimension + 2, BigInteger.ZERO, order);
    }

    /**
     * 初始化加密密钥组（生成可逆矩阵的逆矩阵）
     *
     * @param g      循环群生成元
     * @param n         取自有效域的随机值
     * @param matrix 已生成的随机可逆矩阵
     */
    public EncryptionKey(Element g, Element n, BigIntMatrix matrix) {
        // 根据所给参数是一个随机可逆矩阵，确定该组密钥需要生成该矩阵的逆矩阵，也就是说，这组密钥用于加密数据向量
        ifData = true;

        // 记录属于 Z^*_p 的随机值 n，对生成元 g 进行 n 次运算
        this.n = n;
        number = g.powZn(n).getImmutable().getElementPowPreProcessing();

        // 根据所给矩阵，生成该矩阵的逆矩阵
        keyMatrix = matrix.inverse();
    }

    /**
     * 生成可逆矩阵
     *
     * @param dimensions 维度数
     * @return {@link Matrix} 可逆矩阵
     */
    private Matrix generateInvertibleMatrix(int dimensions) {
        Matrix matrix = null;
        boolean isMatrixInvertible = false;

        while (!isMatrixInvertible) {
            double[][] randomData = new double[dimensions][dimensions];
            for (int i = 0; i < dimensions; i++) {
                for (int j = 0; j < dimensions; j++) {
                    randomData[i][j] = ThreadLocalRandom.current().nextDouble(-10, 10); //每一位生成一个0.1到0.9之间的随机浮点数
                }
            }
            matrix = new Matrix(randomData);
            isMatrixInvertible = isMatrixInvertible(matrix);
        }

        return matrix;
    }

    /**
     * 矩阵是否可逆
     *
     * @param matrix 矩阵
     * @return boolean 若为可逆矩阵，返回true，否则返回false
     */
    private boolean isMatrixInvertible(Matrix matrix) {
        try {
            matrix.inverse();
            return true;
        } catch (RuntimeException e) {
            return false;
        }
    }

    /**
     * 加密向量（第二步）
     *
     * @param vector1 已经经过第一步加密的向量
     * @return {@link Element[]}
     */
    public Element[] encrypt(BigInteger[] vector1, PairingParameter pairingParameter) {
        // 对向量进行二次加密
        BigInteger[] vector2;
        // 根据当前密钥组的性质，分情况对向量进行加密
        if (ifData) {
            // 若当前密钥组用于加密数据向量，加密格式是 keyMatrix^{-1} * vector1
            vector2 = keyMatrix.times4(vector1);
        } else {
            // 若当前密钥组用于加密数据向量，加密格式是 vector1 * keyMatrix
            vector2 = keyMatrix.times(vector1);
        }

        // 对向量进行三次加密，形式为number^vector2
        Element[] vector3 = new Element[vector2.length];
        for (int i = 0; i < vector3.length; i++) {
            vector3[i] = number.powZn(pairingParameter.getFromZp(vector2[i]).getImmutable()).getImmutable();
        }

        return vector3;
    }

    public BigIntMatrix getKeyMatrix() {
        return keyMatrix;
    }

    public ElementPowPreProcessing getNumber() {
        return number;
    }

    public Element getN() {
        return n;
    }
}
